using System;
using System.Runtime.Remoting.Messaging;
using System.Security.Permissions;
using System.Transactions;

namespace MSTestExtensions
{
    /// <remarks />
    public class TestTransactionAspect : TestAspect<TestTransactionAttribute>, IMessageSink, ITestAspect
    {
        #region Fields

        /// <remarks />
        private IMessageSink _nextSink;

        #endregion // Fields

        #region IMessageSink Members

        /// <remarks />
        [SecurityPermission(SecurityAction.LinkDemand, Flags = SecurityPermissionFlag.Infrastructure)]
        public IMessage SyncProcessMessage(IMessage msg)
        {
            if (msg == null)
                throw new ArgumentNullException("msg");
            TestTransactionAttribute testTransactionAttribute = GetAttribute(msg);
            IMessage returnMethod;
            if (testTransactionAttribute != null)
            {
                using (TransactionScope scope = new TransactionScope(testTransactionAttribute.TransactionScopeOption, testTransactionAttribute.TransactionOptions))
                {
                    returnMethod = _nextSink.SyncProcessMessage(msg);
                    if (testTransactionAttribute.Commit)
                    {
                        Console.WriteLine("Commiting Transaction in Test Method.");
                        scope.Complete();
                    }
                }
            }
            else
            {
                returnMethod = _nextSink.SyncProcessMessage(msg);
            }
            return returnMethod;
        }

        /// <remarks />
        [SecurityPermission(SecurityAction.LinkDemand, Flags = SecurityPermissionFlag.Infrastructure)]
        public IMessageCtrl AsyncProcessMessage(IMessage msg, IMessageSink replySink)
        {
            throw new InvalidOperationException();
        }

        /// <remarks />
        public IMessageSink NextSink
        {
            [SecurityPermission(SecurityAction.LinkDemand, Flags = SecurityPermissionFlag.Infrastructure)]
            get { return _nextSink; }
        }

        #endregion // IMessageSink Members

        #region ITestAspect

        /// <remarks />
        public void AddMessageSink(IMessageSink messageSink)
        {
            _nextSink = messageSink;
        }

        #endregion // ITestAspect
    }
}
